import tensorflow as tf
import numpy as np

import util
from lenet import LeNet

if __name__ == '__main__':
    lenet = LeNet()
    
    train_set_loader = util.DataLoader('./train-set')
    test_set_loader = util.DataLoader('./test-set')

    batch_size = 200
    total_batches = train_set_loader.total // batch_size

    validation_x, validation_y = test_set_loader.load(0, test_set_loader.total // 2)
    
    saver = tf.train.Saver()
    sess = tf.Session()
    sess.run(tf.global_variables_initializer())
    max_validation_accuracy = 0
    for epoch in range(50):
        for i in range(total_batches):
            print("Epoch %d: %.2f%%" % (epoch, 100 * i / total_batches), end='\r')
            x, y = train_set_loader.load(i * batch_size, batch_size)
            sess.run(lenet.train_op, feed_dict={
                lenet.input_x_placeholder: x,
                lenet.input_y_placeholder: y
            })
        current_loss, current_accuracy = sess.run(
            [lenet.loss, lenet.accuracy], 
            feed_dict={
                lenet.input_x_placeholder: validation_x,
                lenet.input_y_placeholder: validation_y
            }
        )
        print("* Epoch %d: loss = %f, accuracy = %f" % (epoch, current_loss, current_accuracy))
        if max_validation_accuracy < current_accuracy:
            max_validation_accuracy = current_accuracy
            saver.save(sess, './model/model')
